from multiprocessing import Process, Queue
import os
import time
from lxml import etree

from mzitu.modular.headers import Headers
from mzitu.modular.database import DataBase


class Mzitu(object):
    def __init__(self, all_link, headers):
        self.start_url = all_link
        self.mzitu_request = headers
        self.queue = Queue()

    def save_term(self):
        req = self.mzitu_request.get(url=self.start_url, timeout=30)
        html = etree.HTML(req.text)
        links = html.xpath('//div[@class="all"]//a')
        for link in links:
            term = {
                'link': link.xpath('@href')[0],
                'title': link.text
            }
            self.queue.put(term)

    def get_html(self):
        db = DataBase('girls', 'mzitu').get_collection()
        while not self.queue.empty():
            term = self.queue.get()
            time.sleep(5)
            term_link = term.get('link')
            term_title = term.get('title')
            path = './pic/' + term_title
            path = path.replace('?', '')
            path = path.replace(':', '')
            path = path.rstrip()
            print('正在下载：' + term_title)
            if not db.find_one({'_id': term_link}):
                data = {
                    '_id': term_link,
                    'title': term_title,
                    'pic_details': []
                }
                self.creat_mkdir(path)
                db.save(data)
            self.get_pic(term_url=term_link, file_name=path + '/')

    def get_pic(self, term_url, file_name):
        db = DataBase('girls', 'mzitu').get_collection()
        db_data = db.find_one({'_id': term_url})
        count = len(db_data['pic_details'])
        counter = 3
        while counter > 0:
            req = self.mzitu_request.get(url=term_url, timeout=30)
            html = etree.HTML(req.text)
            max_page = html.xpath('//div[@class="pagenavi"]/a[5]/span/text()')
            if max_page:
                max_page = max_page[0]
                break
            counter -= 1
        if counter < 0:
            return
        for index in range(count+1, int(max_page)+1):
            pic_url = term_url + '/' + str(index)
            counter = 3
            while counter > 0:
                req = self.mzitu_request.get(url=pic_url, timeout=30)
                html = etree.HTML(req.text)
                pic_src = html.xpath('//div[@class="main-image"]//img/@src')
                if pic_src:
                    pic_src = pic_src[0]
                    break
                counter -= 1
            if counter < 0:
                return
            if not db.find_one({'pic_details.pic_link': pic_url}):
                pic_path = self.save_pic(pic_src=pic_src, file_name=file_name)
                _id = {
                    '_id': term_url
                }
                data = {'$push': {'pic_details': {
                    'pic_link': pic_url, 'pic_path': pic_path}}}
                db.update_one(_id, data)

    def save_pic(self, pic_src, file_name):
        req = self.mzitu_request.get(url=pic_src, timeout=30)
        pic_name = os.path.basename(pic_src)
        path = file_name + pic_name
        is_exists = os.path.exists(path)
        if not is_exists:
            with open(path, 'wb') as f:
                f.write(req.content)
            return path
        else:
            return False

    """文件夹创建模块
    args:
        path:所创建文件夹路径及文件夹名
    """
    @staticmethod
    def creat_mkdir(path):
        is_exists = os.path.exists(path)
        if not is_exists:
            try:
                os.makedirs(path)
            except OSError:
                pass
        else:
            # print('该文件夹已存在')
            pass


if __name__ == '__main__':
    start_url = 'http://www.mzitu.com/all'
    mzitu_request = Headers()
    mzitu = Mzitu(start_url, mzitu_request)
    mzitu.save_term()
    for i in range(10):
        p = Process(target=mzitu.get_html)
        p.start()
